#!/usr/bin/env python3
import numpy as np 
from kitti_util import *

def compute_3d_box_cam2(h, w, l, x, y, z, yaw, calib):
    """
    Return: 3xn in cam2 coordinate
    """
    calib_data = calib.read_calib_from_video('/home/jony/data/kitti/RawData/2011_09_26/')
    # 获取投影矩阵 
    P = np.reshape(calib_data['P2'], (3, 4))
    # 旋转矩阵
    R = np.array([[np.cos(yaw), 0, np.sin(yaw)], [0, 1, 0], [-np.sin(yaw), 0, np.cos(yaw)]])
    # 3d bounding box corners
    x_corners = [l/2, l/2, -l/2, -l/2, l/2, l/2, -l/2, -l/2]
    y_corners = [0, 0, 0, 0, -h, -h, -h, -h]
    z_corners = [w/2, -w/2, -w/2, w/2, w/2, -w/2, -w/2, w/2]
    corners_3d_cam2 = np.dot(R, np.vstack([x_corners, y_corners, z_corners]))
    corners_3d_cam2 += np.vstack([x, y, z])
    # only draw 3d bounding box for objs in front of the camera
    # if np.any(corners_3d_cam2[2,:] < 0.1):
    #     corners_2d = None
    #     return corners_2d, corners_3d_cam2
    # project the 3d bounding box into the image plane
    corners_2d = project_to_image(np.transpose(corners_3d_cam2), P)
    return corners_2d, corners_3d_cam2

def get_2d_boxes(corners_2d):
    """
    根据2D角点获取最小边界矩形
    corners_2d: (8, 2) array, 每个元素表示二维角点的 (x, y) 坐标
    返回值:
    bbox: (4,) array, 边界矩形的左上角和右下角点的坐标 (x_min, y_min, x_max, y_max)
    """
    x_min = np.min(corners_2d[:, 0])
    y_min = np.min(corners_2d[:, 1])
    x_max = np.max(corners_2d[:, 0])
    y_max = np.max(corners_2d[:, 1])
    
    bbox = np.array([x_min, y_min, x_max, y_max], dtype=int)
    return bbox

def convert_world_to_camera_coords(tx, ty, tz, calibration):
    world_coords = np.array([tx, ty, tz, 1.0]).reshape(4, 1)
    camera_coords = np.dot(calibration.V2C, world_coords)
    return camera_coords[0, 0], camera_coords[1, 0], camera_coords[2, 0]

# 计算3D旋转矩阵 
def wrap_to_pi(angle): 
    return (angle + np.pi) % (2 * np.pi) - np.pi

def rz_to_ry(rz):
    angle = -rz - np.pi / 2

    if angle >= np.pi:
        angle -= np.pi
    if angle < -np.pi:
        angle = 2 * np.pi + angle
    return angle

# def rz_to_ry(rz):
#     angle = -rz + np.pi / 2

#     if angle >= np.pi:
#         angle -= np.pi
#     elif angle < -np.pi:
#         angle += 2 * np.pi
#     return angle


def distance_point_to_segment(P, A, B):
    """
    Calculates the min distance of a point P to a segment AB in 3D space.
    Returns the distance and the point Q on segment AB where the min distance is reached.
    """
    AP = P - A
    BP = P - B
    AB = B - A
    if np.dot(AB, AP) >= 0 and np.dot(-AB, BP) >= 0:
        return np.linalg.norm(np.cross(AP, AB)) / np.linalg.norm(AB), np.dot(AP, AB) / np.dot(AB, AB) * AB + A
    d_PA = np.linalg.norm(AP)
    d_PB = np.linalg.norm(BP)
    if d_PA < d_PB:
        return d_PA, A
    return d_PB, B

def min_distance_cuboids(cub1, cub2):
    """
    Computes the minimum distance between two non-overlapping cuboids in 3D space.
    They are projected to 3D space and the minimum distance between the two cuboids is returned.
    """
    minD = 1e5  # 设定一个较大的值
    for i in range(8):
        for j in range(8):
            d, Q = distance_point_to_segment(cub1[i], cub2[j], cub2[(j + 1) % 8])
            if d < minD:
                minD = d
                minP = cub1[i]
                minQ = Q
    # 更换方框顶点
    for i in range(8):
        for j in range(8):
            d, Q = distance_point_to_segment(cub2[i], cub1[j], cub1[(j + 1) % 8])
            if d < minD:
                minD = d
                minP = cub2[i]
                minQ = Q
    return minP, minQ, minD


# 计算线段0-2和1-3的交点（3D检测框底面交点）
def cross_point_3d0123(P0, P1, P2, P3):
    d_02 = P2 - P0
    d_13 = P3 - P1

    # 解线性方程
    A = np.array([d_02[:2], -d_13[:2]]).T
    b = P1[:2] - P0[:2]
    t, s = np.linalg.solve(A, b)